import numpy as np
from .tac_frame import TacFrame
from typing import List, Dict, Tuple
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


def extract_tacdata_from_frame(frame: Dict, type: str = "PDF"):
    '''
    Extract Tac3D data from a TacFrame dictionary.
    Args:
        frame (Dict): TacFrame dictionary.
        type (str, optional): Type of data to extract. Options are "PDF", "resultant", "PDFr", "PDF_tight" or "all". Defaults to "PDF".
    Returns:
        Depending on the type, returns:
        - "PDF": 3D Positions, 3D Displacements, 3D Forces
        - "resultant": 3D Resultant Force, 3D Resultant Moment
        - "PDFr": 3D Positions, 3D Displacements, 3D Forces, 3D Resultant Force, 3D Resultant Moment
        - "PDF_tight": an (400,9) np.ndarray with 3D Positions, 3D Displacements, 3D Forces
        - "all": TacFrame object
    '''
    tacframe = TacFrame(frame)
    if type == "PDF":
        return tacframe.P, tacframe.D, tacframe.F
    elif type == "resultant":
        return tacframe.Fr, tacframe.Mr
    elif type == "PDFr":
        return tacframe.P, tacframe.D, tacframe.F, tacframe.Fr, tacframe.Mr
    elif type == "PDF_tight":
        return np.concatenate((tacframe.P, tacframe.D, tacframe.F), axis=1)
    elif type == "all":
        return tacframe
    else:
        raise ValueError("type must be 'PDF', 'resultant', 'PDFr' or 'all'")

def zip_PDF(P: np.ndarray, D: np.ndarray, F: np.ndarray) -> np.ndarray:
    '''
    Zip 3D Positions, 3D Displacements, and 3D Forces into a single array.
    Args:
        P (np.ndarray): 3D Positions of shape (N, 400, 3) or (N, 20,20,3).
        D (np.ndarray): 3D Displacements of shape (N, 400, 3) or (N, 20,20,3).
        F (np.ndarray): 3D Forces of shape (N, 400, 3) or (N, 20,20,3).
    Returns:
        np.ndarray: A zipped array of shape (N, 400, 9) or (N, 20,20,3) containing [P, D, F].
    '''
    return np.concatenate((P, D, F), axis=-1)

def unzip_PDF(PDF: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    '''
    Unzip a zipped array of 3D Positions, 3D Displacements, and 3D Forces.
    Args:
        PDF (np.ndarray): A zipped array of shape (N, 400, 9) or (N, 20,20,9).
    Returns:
        tuple: A tuple containing:
            - P (np.ndarray): 3D Positions of shape (N, 400, 3) or (N, 20,20,3).
            - D (np.ndarray): 3D Displacements of shape (N, 400, 3) or (N, 20,20,3).
            - F (np.ndarray): 3D Forces of shape (N, 400, 3) or (N, 20,20,3).
    '''
    return (PDF[..., :3], PDF[..., 3:6], PDF[..., 6:])

